/*
 * Copyright(c) YouYuan 2018.
 * https://gitee.com/you-yuan
 */

package com.yuan.ocr.kit;

import com.yuan.ocr.common.OcrConst;
import com.yuan.ocr.tess4j.Tess4jOCR;
import com.yuan.util.file.FileUtil;
import com.yuan.util.image.ImageUtil;
import com.yuan.util.string.StrUtil;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 识别html中的图片上的文字
 * 目前仅解析img标签上src值为图片的base64编码
 * @author YouYuan
 * @version 1.0
 * @Email xiyousuiyuan@163.com
 * @Date 2018/10/28 17:56
 */
public class HtmlImgOCR {

    /**
     * 识别html中的图片上的文字
     * 目前仅解析img标签上src值为图片的base64编码
     * @param htmlFilePath html文件路径
     * @param charset html文件字符编码
     * @param tempPath 图片临时保存路径
     * @return
     * @throws Exception
     */
    public static List<String> doOcr(String htmlFilePath, String charset, String tempPath) throws Exception {
        List<String> result = new ArrayList<String>();
        Document document = Jsoup.parse(new File(htmlFilePath), charset);
        Elements imgs = document.getElementsByTag("img");
        Iterator<Element> iterator = imgs.iterator();
        int size = imgs.size();
        System.out.println("[Info] HtmlImgOCR开始处理，共 " + size + " 张图片");
        String imgPath = FileUtil.formatPath(tempPath + "/");
        int index = 1;
        while (iterator.hasNext()) {
            System.out.println("[Info] 正在处理第 " + index + "/" + size + " 张图片...");
            Element e = iterator.next();
            String src = e.attr("src");
            byte[] bytes = StrUtil.base64ToByteArray(src);
            String img = imgPath + index + ".jpg";
            FileOutputStream fos = new FileOutputStream(img);
            fos.write(bytes);
            fos.close();
            //因为base64转换的图片可能不识别，进行处理后再执行文字识别操作
            BufferedImage bufferedImage = ImageUtil.zoomImage(img, 0);
            ImageUtil.saveImage(bufferedImage, imgPath, index + "p" + ".jpg");
            String imgp = imgPath + index + "p" + ".jpg";
            String ocr = Tess4jOCR.doOCR(imgp, OcrConst.LANGUAGE_CHI_SIM);
            result.add(ocr);
            index++;
            new File(img).delete();
        }
        System.out.println("[Info] HtmlImgOCR图片全部处理完成");
        return result;
    }

}
